package utils

import "errors"

// 多叉树（类似前缀树）

type CTreeNode[K any, V any] struct {
	Key      K
	Value    V
	Children []*CTreeNode[K, V]
}

type CTree[K any, V any] struct {
	rw         rwMutex
	root       *CTreeNode[K, V]
	comparator Comparator[K]
}

// a<b 则返回小于0 ， a==b则返回0 ， a>b则返回大于0
type Comparator[K any] func(a, b K) int

func NewCommonTree[K any, V any](c Comparator[K]) *CTree[K, V] {
	return &CTree[K, V]{
		comparator: c,
		root:       &CTreeNode[K, V]{},
	}
}

func (t *CTree[K, V]) WithLock(b ...bool) {
	t.rw = rwCreate(b...)
}

// 泛查找
// 通过key查找树的节点，keys必须按照树的顺序
// 比如：1 下有 2 3 4 三个节点， 2 下有 5 6节点，传入 1 2 7 ,会返回2节点
func (t *CTree[K, V]) GetNodeNormal(keys ...K) (*CTreeNode[K, V], bool) {
	t.rw.RLock()
	defer t.rw.RUnlock()
	return t.getNodeNormal(keys...)
}
func (t *CTree[K, V]) getNodeNormal(keys ...K) (*CTreeNode[K, V], bool) {
	node := t.root
	isRoot := true
	for _, k := range keys {
		if n, ok := node.findNodeByKey(k, t.comparator); ok {
			isRoot = false
			node = n
		} else {
			break
		}
	}

	if isRoot {
		return nil, false
	}
	return node, true
}

// 严格查找
// 通过key查找树的节点，keys必须按照树的顺序
// 比如：1 下有 2 3 4 三个节点， 2 下有 5 6节点，如果要获取5节点，则必须按照 1 2 5的顺序传参，否则查找不到
func (t *CTree[K, V]) GetNodeStrict(keys ...K) (*CTreeNode[K, V], bool) {
	t.rw.RLock()
	defer t.rw.RUnlock()
	return t.getNodeStrict(keys...)
}
func (t *CTree[K, V]) getNodeStrict(keys ...K) (*CTreeNode[K, V], bool) {
	node := t.root
	for _, k := range keys {
		if n, ok := node.findNodeByKey(k, t.comparator); ok {
			node = n
		} else {
			return nil, false
		}
	}
	return node, true
}

// 获取节点，如果没有则创建
func (t *CTree[K, V]) GetOrCreateNode(keys ...K) *CTreeNode[K, V] {
	t.rw.Lock()
	defer t.rw.Unlock()

	curr := t.root
	for _, k := range keys {
		if n, ok := curr.findNodeByKey(k, t.comparator); !ok {
			newNode := &CTreeNode[K, V]{
				Key: k,
			}
			curr.addChild(newNode)
			curr = newNode
		} else {
			curr = n
		}
	}
	return curr
}

// 修改节点的值，如果不存在则插入
func (t *CTree[K, V]) SetOrCreateNode(val V, keys ...K) error {
	t.rw.Lock()
	defer t.rw.Unlock()

	curr := t.root
	for _, k := range keys {
		if n, ok := curr.findNodeByKey(k, t.comparator); !ok {
			newNode := &CTreeNode[K, V]{
				Key: k,
			}
			curr.addChild(newNode)
			curr = newNode
		} else {
			curr = n
		}
	}
	curr.Value = val
	return nil
}

// 将node节点插入到keys的子节点中
func (t *CTree[K, V]) InsertNode(node *CTreeNode[K, V], keys ...K) error {
	t.rw.Lock()
	defer t.rw.Unlock()

	curr := t.root
	for _, k := range keys {
		if n, ok := curr.findNodeByKey(k, t.comparator); !ok {
			newNode := &CTreeNode[K, V]{
				Key: k,
			}
			curr.addChild(newNode)
			curr = newNode
		} else {
			curr = n
		}
	}
	curr.addChild(node)
	return nil
}

// 删除节点
func (t *CTree[K, V]) DeleteNode(keys ...K) error {
	t.rw.Lock()
	defer t.rw.Unlock()

	if len(keys) <= 0 {
		return errors.New("invalid keys")
	}
	n, ok := t.getNodeStrict(keys[:len(keys)-1]...)
	if !ok {
		return nil
	}
	ind, found := n.findKeyInChildren(keys[len(keys)-1], t.comparator)
	if !found {
		return nil
	}

	// 删除对应的子结点
	//n.Children = append(n.Children[:ind], n.Children[ind+1:]...)
	copy(n.Children[:ind], n.Children[ind+1:])
	n.Children = n.Children[:len(n.Children)-1]

	return nil
}

func (n *CTreeNode[K, V]) addChild(child *CTreeNode[K, V]) {
	n.Children = append(n.Children, child)
}

func (n *CTreeNode[K, V]) findNodeByKey(key K, compare Comparator[K]) (*CTreeNode[K, V], bool) {
	if compare(n.Key, key) == 0 {
		return n, true
	}
	if ind, ok := n.findKeyInChildren(key, compare); ok {
		return n.Children[ind], true
	}
	return nil, false
}

func (n *CTreeNode[K, V]) findKeyInChildren(key K, c Comparator[K]) (int, bool) {
	for i, child := range n.Children {
		if c(child.Key, key) == 0 {
			return i, true
		}
	}
	return 0, false
}

func (n *CTreeNode[K, V]) GetChildCount() int {
	return len(n.Children)
}

// 获取所有叶子节点
func (t *CTree[K, V]) GetAllLeafNodes() []*CTreeNode[K, V] {
	t.rw.RLock()
	defer t.rw.RUnlock()
	return t.getLeafNodesByNum(t.root, 0)
}

// 获取指定数量的叶子节点
func (t *CTree[K, V]) GetLeafNodesByNum(count int) []*CTreeNode[K, V] {
	t.rw.RLock()
	defer t.rw.RUnlock()
	return t.getLeafNodesByNum(t.root, count)
}

// 获取某节点下 指定数量 的叶子节点
func (t *CTree[K, V]) GetLeafByKey(count int, keys ...K) ([]*CTreeNode[K, V], bool) {
	t.rw.RLock()
	defer t.rw.RUnlock()
	n, ok := t.getNodeStrict(keys...)
	if !ok {
		return nil, false
	}

	res := t.getLeafNodesByNum(n, count)
	return res, len(res) != 0
}

func (t *CTree[K, V]) getLeafNodesByNum(n *CTreeNode[K, V], num int) []*CTreeNode[K, V] {
	leafNodes := make([]*CTreeNode[K, V], 0, num)

	var traverse func(node *CTreeNode[K, V])
	traverse = func(node *CTreeNode[K, V]) {
		if len(node.Children) == 0 {
			leafNodes = append(leafNodes, node)
			if num != 0 {
				if len(leafNodes) == num {
					return
				}
			}
		} else {
			for _, child := range node.Children {
				traverse(child)
			}
		}
	}

	traverse(n)
	return leafNodes
}
